import os
import sys

import torch
import numpy as np
import potpourri3d as pp3d

sys.path.append(os.path.join("D:/_Code_library/diffusion-net/src"))  # add the path to the DiffusionNet src
import diffusion_net
from diffusion_net.utils import toNP


def compute_variables(base_dir, fnames, k_eig):
    # Build a small dataset with 2 meshes for evaluation
    # Initialize
    verts_list = []
    faces_list = []
    normals_list = []

    for i in range(len(fnames)):

        # load mesh
        mesh_fullpath = os.path.join(base_dir, fnames[i])
        verts, faces = pp3d.read_mesh(mesh_fullpath)
        normals = None
        # convert to torch
        verts = torch.tensor(np.ascontiguousarray(verts)).float()
        faces = torch.tensor(np.ascontiguousarray(faces))
        # center and scale
        verts = diffusion_net.geometry.normalize_positions(verts, method='bbox')

        verts_list.append(verts)
        faces_list.append(faces)
        normals_list.append(normals)

    # Compute operators
    frames_list, massvec_list, L_list, evals_list, evecs_list, gradX_list, gradY_list = diffusion_net.geometry.get_all_operators(verts_list, faces_list, normals=normals_list, k_eig=k_eig)

    # Evecs transpose
    evecs_trans_list = []
    for i in range(len(verts_list)):
        evecs_trans = evecs_list[i].t() @ torch.diag(massvec_list[i])
        evecs_trans_list.append(evecs_trans)

    return verts_list, faces_list, frames_list, massvec_list, L_list, evals_list, evecs_list, gradX_list, gradY_list, evecs_trans_list

def compute_variables_area(base_dir, fnames, k_eig):
    # Build a small dataset with 2 meshes for evaluation
    # Initialize
    verts_list = []
    faces_list = []
    normals_list = []

    for i in range(len(fnames)):

        # load mesh
        mesh_fullpath = os.path.join(base_dir, fnames[i])
        verts, faces = pp3d.read_mesh(mesh_fullpath)
        normals = None

        # center and scale - per-area
        verts -= np.mean(verts, axis=0)
        vertex_areas = pp3d.vertex_areas(verts, faces)
        surface_area = vertex_areas.sum()
        verts /= np.sqrt(surface_area) # Area-normalized

        # convert to torch
        verts = torch.tensor(np.ascontiguousarray(verts)).float()
        faces = torch.tensor(np.ascontiguousarray(faces))
        # center and scale
        # verts = diffusion_net.geometry.normalize_positions(verts, method='bbox')

        verts_list.append(verts)
        faces_list.append(faces)
        normals_list.append(normals)

    # Compute operators
    frames_list, massvec_list, L_list, evals_list, evecs_list, gradX_list, gradY_list = diffusion_net.geometry.get_all_operators(verts_list, faces_list, normals=normals_list, k_eig=k_eig)

    # Evecs transpose
    evecs_trans_list = []
    for i in range(len(verts_list)):
        evecs_trans = evecs_list[i].t() @ torch.diag(massvec_list[i])
        evecs_trans_list.append(evecs_trans)

    return verts_list, faces_list, frames_list, massvec_list, L_list, evals_list, evecs_list, gradX_list, gradY_list, evecs_trans_list
